import pandas as pd
import os
import matplotlib.pyplot as plt
import re


def main():
    if not os.path.exists('images'):
        os.makedirs('images')
    df = pd.read_csv("csv/dalian_weather_2022-2024_full.csv")

    # 数据清洗
    df['日期'] = pd.to_datetime(df['日期'])
    df['月份'] = df['日期'].dt.month
    df['年份'] = df['日期'].dt.year

    # 标准化风力描述
    def standardize_wind_description(wind_str):
        wind_str = wind_str.replace(' ', '')
        wind_str = re.sub(r'(无持续风向|[东南西北]{1,2}风)', '', wind_str)
        return wind_str

    df['白天风力等级'] = df['白天风力风向'].apply(standardize_wind_description)
    df['夜间风力等级'] = df['夜间风力风向'].apply(standardize_wind_description)

    def count_wind_levels_avg(data, wind_column):
        yearly_counts = data.groupby(['年份', '月份', wind_column]).size().unstack().fillna(0)
        monthly_avg = yearly_counts.groupby('月份').mean()
        return monthly_avg

    day_wind_avg = count_wind_levels_avg(df, '白天风力等级')
    night_wind_avg = count_wind_levels_avg(df, '夜间风力等级')
    combined_wind_avg = day_wind_avg.add(night_wind_avg)

    # 设置中文显示
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False

    # 定义月份名称
    month_names = ['1月', '2月', '3月', '4月', '5月', '6月', '7月', '8月', '9月', '10月', '11月', '12月']

    # 获取所有风力等级并排序
    all_levels = sorted(set(day_wind_avg.columns) | set(night_wind_avg.columns),
                        key=lambda x: int(re.search(r'\d+', x).group(0)) if re.search(r'\d+', x) else 0)

    # 为所有风力等级创建颜色映射
    def create_wind_colormap(levels):
        cmap = plt.get_cmap('Set3')
        colors = {}
        level_count = len(levels)
        for i in range(level_count):
            ratio = float(i) / level_count
            color_value = cmap(ratio)
            current_level = levels[i]
            colors[current_level] = color_value
        return colors

    color_map = create_wind_colormap(all_levels)

    # 绘制饼图
    def plot_wind_pie(data, title, filename):
        plt.figure(figsize=(16, 14))
        plt.suptitle(title, fontsize=22, y=0.98)

        for month in range(1, 13):
            plt.subplot(4, 3, month)
            if month in data.index:
                month_data = data.loc[month]
                month_data = month_data[month_data > 0]

                if month_data.empty:
                    plt.text(0.5, 0.5, f'{month_names[month - 1]}\n无数据',
                             ha='center', va='center', fontsize=12)
                    plt.axis('off')
                    continue

                # 为当前月份的风力等级分配颜色
                pie_colors = [color_map[level] for level in month_data.index]

                # 计算百分比
                total_days = month_data.sum()
                percentages = month_data / total_days * 100

                # 设置标签 - 只显示大于5%的部分
                labels = []
                for level, pct in zip(month_data.index, percentages):
                    if pct > 5:
                        labels.append(f'{level}\n({pct:.1f}%)')
                    else:
                        labels.append('')

                # 绘制饼图
                plt.pie(month_data,
                        labels=labels,
                        colors=pie_colors,
                        startangle=90,
                        wedgeprops={'edgecolor': 'w', 'linewidth': 1},
                        textprops={'fontsize': 9, 'fontweight': 'bold'})

                plt.title(f'{month_names[month - 1]}', fontsize=14, pad=15)
            else:
                plt.text(0.5, 0.5, f'{month_names[month - 1]}\n无数据',
                         ha='center', va='center', fontsize=12)
                plt.axis('off')

        # 添加图例
        handles = [plt.Rectangle((0, 0), 1, 1, color=color_map[level]) for level in all_levels]
        plt.figlegend(handles, all_levels,
                      loc='lower center',
                      ncol=min(6, len(all_levels)),
                      title='风力等级',
                      fontsize=11,
                      title_fontsize=12,
                      bbox_to_anchor=(0.5, 0.01))

        plt.tight_layout(rect=(0, 0.05, 1, 0.95))
        plt.savefig(f'../weather/images/{filename}', dpi=300, bbox_inches='tight')
        plt.show()

    # 1. 绘制白天风力饼图
    plot_wind_pie(day_wind_avg,
                  '大连近三年各月白天风力等级分布 (2022-2024)',
                  'dalian_day_wind_pie_avg.png')

    # 2. 绘制夜间风力饼图
    plot_wind_pie(night_wind_avg,
                  '大连近三年各月夜间风力等级分布 (2022-2024)',
                  'dalian_night_wind_pie_avg.png')

    # 3. 绘制综合风力饼图
    plot_wind_pie(combined_wind_avg,
                  '大连近三年各月综合风力等级分布 (2022-2024)',
                  'dalian_combined_wind_pie_avg.png')

    print("饼图已保存到../weather/images/目录")


if __name__ == '__main__':
    main()